/*
 * This code implements stVSSIM VQA as suggested by Anush Moorthy and A.C Bovik
 * in "Efficient Motion Weighted Spatio-Temporal Video SSIM Index"
 * http://live.ece.utexas.edu/publications/2010/moorthy_spie_jan10.pdf
 *
 * Please cite the following paper in any published work if you use this software.
 * Anus K. Moorthy, Alan C. Bovik,"Efficient Motion Weighted Spatio-Temporal
 * Video SSIM Index", SPIE Conference on Human Vision and Electronic Imaging,
 * San Jose, CA, January 17-21, 2010
 *
 * File: calcMSSIM.cpp
 *
 * This module calculates the MC-SSIM for distorted image with reference to
 * the original image
 */

#include "vssim.h"

extern vssim_global_t global;
extern vssim_video_t refVideo;
extern vssim_video_t distortedVideo;
extern vssim_FrScore_t *Tframe;

vssim_filters_t Filters;
vssim_ImgStats_t Ref_ImgStats;
vssim_ImgStats_t Dist_ImgStats;
vssim_FilStats_t CovarSig;

vssim_FilStats_t SSIM3D;

extern uint32_t findMax(IplImage *img);
extern float calcPSSIM(IplImage *ssim_map);
extern void calcMotionVectors(FILE *ref, int cntframe, vssim_size_t size, vssim_ImgComp_t comp,
								  vssim_mv_t *mv, int block_size, int num_blocks);

/*
 * Creates H filter, y = 0
 */
static void createHFilter(int dim)
{
	int i, j;
	double kernel[121] = {0.0};

	// H filter == {1} at the horizontal gradient
	int value = 1;
	for (j = 0; j < dim; j++)
	{
		for (i = 0; i < dim; i++)
		{
			if (j == 5)
				kernel[j * dim + i] = value;
		}
	}
	cvSetData(Filters.H, kernel, dim*4);

	return;
}

/*
 * Creates the V filter, x = 0
 */
static void createVFilter(int dim)
{
	int i, j;
	double kernel[121] = {0.0};

	//V filter = {1} for 6th col = transpose of H filter
	int value = 1;
	for (i = 0; i < dim; i++)
	{
		for (j = 0; j < dim; j++)
		{
			if (i == 5)
				kernel[j*dim + i] = value;
		}
	}
	cvSetData(Filters.V, kernel, dim*4);

	return;
}

/*
 * Creates the Left filter, x = -y
 */
static void createLFilter(int dim)
{
	double 	kernel[] = { 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
  						 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0,
						 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0,
						 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,
						 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0,
						 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,
						 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0,
						 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0,
						 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0,
						 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0,
						 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1 };

	cvSetData(Filters.L, kernel, dim*4);

	return;
}

/*
 * Creates the right filter x = y
 */
static void createRFilter(int dim)
{
	double 	kernel[] = { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1,
						 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0,
						 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0,
						 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0,
						 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0,
						 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0,
						 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0,
						 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0,
						 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0,
						 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0,
						 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};

	cvSetData(Filters.R, kernel, dim*4);

	return;
}

/*
 * This function creates Horizontal, Vertical, Left and Right filters
 */
static void createFilters()
{
	int dim = 11;

	//Create the matrix headers
	Filters.H = cvCreateMatHeader(dim, dim, CV_32FC1);
	Filters.V = cvCreateMatHeader(dim, dim, CV_32FC1);
	Filters.L = cvCreateMatHeader(dim, dim, CV_32FC1);
	Filters.R = cvCreateMatHeader(dim, dim, CV_32FC1);

	// call the create filter function for the respective filter to be created
	createHFilter(dim);
	createVFilter(dim);
	createLFilter(dim);
	createRFilter(dim);

	return;
}

/*
 * Function filters the image passed to it by the kernel and normalizes since
 * the H, V, L and R kernels are not normalized
 */
static void filterImage(const IplImage *FilImg, CvMat *kernel, IplImage *Mu)
{

	IplImage *temp;

	// Apply filter
	temp = cvCreateImage(cvSize(FilImg->width, FilImg->height), IPL_DEPTH_32F, 1);
	cvFilter2D(FilImg, temp, kernel);
	cvConvertScale(temp, Mu, ((float)1/(float)121));
	cvReleaseImage(&temp);

}

/*
 * This function reads the temporal chunk of image, 16 prev and 16 future images
 * for chunk volume of 33 images. It filters every image using the 4 filters
 * and averages them over the 33 images
 */
static void calcMu(FILE *ref, FILE *dist, int cntframe, vssim_size_t size, vssim_ImgComp_t comp)
{
	int i;
	IplImage *temp;

	/* Get 33 frames, 16 previous and 16 future from the file and add them up
	 * separately for Y, U and V */
	IplImage *ref_frame = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_8U, 1);
	IplImage *ref_float =cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *ref_h = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *ref_v = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *ref_l = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *ref_r = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	IplImage *dist_frame = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_8U, 1);
	IplImage *dist_float =cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *dist_h = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *dist_v = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *dist_l = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *dist_r = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	/* offset is used for reading the respective component from the file */
	long offset;
	switch (comp)
	{
		case Y:
			offset = 0;
			break;
		case U:
			offset = global.size.width * HEIGHT_Y(&global);
			break;
		case V:
			offset = global.size.width *(HEIGHT_Y(&global)+HEIGHT_U(&global));
			break;
	}

	/* read through the file and get prev 16 and future 16 frames */
	for (i = (cntframe-16); i <= (cntframe+16); i++)
	{
		/* go to the ith frame in the file and read the buffer */
		fseek(ref, ((i*global.size.width*FRAME_HEIGHT(&global))+offset), SEEK_SET);
		fread(ref_frame->imageData, sizeof(char),
				(size.width * size.height),
				ref);
		temp = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
		cvConvert(ref_frame, temp);
		cvResize(temp, ref_float);
		cvReleaseImage(&temp);

		fseek(dist, ((i*global.size.width*FRAME_HEIGHT(&global))+offset), SEEK_SET);
		fread(dist_frame->imageData, sizeof(char),
				(size.width * size.height),
				dist);
		temp = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
		cvConvert(dist_frame, temp);
		cvResize(temp, dist_float);
		cvReleaseImage(&temp);

		/* filter the image and accumulate */
		IplImage *ref_filt = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
		IplImage *dist_filt = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

		filterImage(ref_float, Filters.H, ref_filt);
		cvAdd(ref_filt, ref_h, ref_h);

		filterImage(ref_float, Filters.V, ref_filt);
		cvAdd(ref_filt, ref_v, ref_v);

		filterImage(ref_float, Filters.L, ref_filt);
		cvAdd(ref_filt, ref_l, ref_l);

		filterImage(ref_float, Filters.R, ref_filt);
		cvAdd(ref_filt, ref_r, ref_r);

		filterImage(dist_float, Filters.H, dist_filt);
		cvAdd(dist_filt, dist_h, dist_h);

		filterImage(dist_float, Filters.V, dist_filt);
		cvAdd(dist_filt, dist_v, dist_v);

		filterImage(dist_float, Filters.L, dist_filt);
		cvAdd(dist_filt, dist_l, dist_l);

		filterImage(dist_float, Filters.R, dist_filt);
		cvAdd(dist_filt, dist_r, dist_r);

		cvReleaseImage(&ref_filt);
		cvReleaseImage(&dist_filt);

	}

	/* average over the chunk volume, as the paper indicates a chunk of
	 * 33 frames is used */
	cvConvertScale(ref_h, Ref_ImgStats.Mu.H, (1.0/33.0), 0.0);
	cvConvertScale(ref_v, Ref_ImgStats.Mu.V, (1.0/33.0), 0.0);
	cvConvertScale(ref_l, Ref_ImgStats.Mu.L, (1.0/33.0), 0.0);
	cvConvertScale(ref_r, Ref_ImgStats.Mu.R, (1.0/33.0), 0.0);

	cvConvertScale(dist_h, Dist_ImgStats.Mu.H, (1.0/33.0), 0.0);
	cvConvertScale(dist_v, Dist_ImgStats.Mu.H, (1.0/33.0), 0.0);
	cvConvertScale(dist_l, Dist_ImgStats.Mu.H, (1.0/33.0), 0.0);
	cvConvertScale(dist_r, Dist_ImgStats.Mu.H, (1.0/33.0), 0.0);

	/* Clean up */
	cvReleaseImage(&ref_frame);
	cvReleaseImage(&ref_float);
	cvReleaseImage(&ref_h);
	cvReleaseImage(&ref_v);
	cvReleaseImage(&ref_l);
	cvReleaseImage(&ref_r);

	cvReleaseImage(&dist_frame);
	cvReleaseImage(&dist_float);
	cvReleaseImage(&dist_h);
	cvReleaseImage(&dist_v);
	cvReleaseImage(&dist_l);
	cvReleaseImage(&dist_r);

	return;
}

/*
 * This function returns the accumulated difference for calculation of sigma
 */
static void getAvg(const IplImage *Img, const vssim_FilStats_t *Mu, vssim_size_t size,
					vssim_FilStats_t *filImg)
{
	IplImage *sub, *squared, *filH, *filV, *filL, *filR;
	sub = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	squared = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	filH = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	filV = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	filL = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	filR = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	// Subtract Mu from image
	cvSub(Img, Mu->H, sub);
	// Square the difference
	cvPow(sub, squared, 2);
	//weight the image using one of the filters
	filterImage(squared, Filters.H, filH);
	// accumulate
	cvAdd(filH, filImg->H, filImg->H);

	cvSub(Img, Mu->V, sub);
	cvPow(sub, squared, 2);
	filterImage(squared, Filters.V, filV);
	cvAdd(filH, filImg->V, filImg->V);

	cvSub(Img, Mu->L, sub);
	cvPow(sub, squared, 2);
	filterImage(squared, Filters.L, filL);
	cvAdd(filH, filImg->L, filImg->L);

	cvSub(Img, Mu->R, sub);
	cvPow(sub, squared, 2);
	filterImage(squared, Filters.R, filR);
	cvAdd(filH, filImg->R, filImg->R);

	cvReleaseImage(&sub);
	cvReleaseImage(&squared);
	cvReleaseImage(&filH);
	cvReleaseImage(&filV);
	cvReleaseImage(&filL);
	cvReleaseImage(&filR);

	return ;
}

/*
 * like the calcMu(), this function calculates the sigma for the video chunk
 * by filtering the difference and averaging it over the 33 video frames
 */
static void calcSigma(FILE *ref, FILE *dist, int cntframe, vssim_size_t size, vssim_ImgComp_t comp)
{
	int i;
	IplImage *temp;

	/* Get 33 frames, 16 previous and 16 future from the file and add them up
	 * separately for Y, U and V */
	IplImage *ref_frame = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_8U, 1);
	IplImage *ref_float =cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	vssim_FilStats_t ref_temp;
	ref_temp.H = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	ref_temp.V = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	ref_temp.L = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	ref_temp.R = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	IplImage *dist_frame = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_8U, 1);
	IplImage *dist_float =cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	vssim_FilStats_t dist_temp;
	dist_temp.H = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	dist_temp.V = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	dist_temp.L = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	dist_temp.R = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	long offset;
	switch (comp)
	{
		case Y:
			offset = 0;
			break;
		case U:
			offset = global.size.width * HEIGHT_Y(&global);
			break;
		case V:
			offset = global.size.width *(HEIGHT_Y(&global)+HEIGHT_U(&global));
			break;
	}

	for (i = (cntframe-16); i <= (cntframe+16); i++)
	{
		/* go to the ith frame in the file and read the buffer */
		fseek(ref, ((i*global.size.width*FRAME_HEIGHT(&global))+offset), SEEK_SET);
		fread(ref_frame->imageData, sizeof(char),
				(size.width * size.height),
				ref);
		temp = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
		cvConvert(ref_frame, temp);
		cvResize(temp, ref_float);
		cvReleaseImage(&temp);

		fseek(dist, ((i*global.size.width*FRAME_HEIGHT(&global))+offset), SEEK_SET);
		fread(dist_frame->imageData, sizeof(char),
				(size.width * size.height),
				dist);
		temp = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
		cvConvert(dist_frame, temp);
		cvResize(temp, dist_float);
		cvReleaseImage(&temp);

		/* Subtract the image from the corresponding Mu of the image, square
		 * the result and accumulate */
		getAvg(ref_float, &Ref_ImgStats.Mu, size, &ref_temp);

		getAvg(dist_float, &Dist_ImgStats.Mu, size, &dist_temp);

	}

	//div the accumulated value by # of images to get the average
	cvConvertScale(ref_temp.H, Ref_ImgStats.Sigma.H, ((float)1/(float)33), 0.0);
	cvConvertScale(ref_temp.V, Ref_ImgStats.Sigma.V, ((float)1/(float)33), 0.0);
	cvConvertScale(ref_temp.L, Ref_ImgStats.Sigma.L, ((float)1/(float)33), 0.0);
	cvConvertScale(ref_temp.R, Ref_ImgStats.Sigma.R, ((float)1/(float)33), 0.0);

	cvConvertScale(dist_temp.H, Dist_ImgStats.Sigma.H, ((float)1/(float)33), 0.0);
	cvConvertScale(dist_temp.V, Dist_ImgStats.Sigma.V, ((float)1/(float)33), 0.0);
	cvConvertScale(dist_temp.L, Dist_ImgStats.Sigma.L, ((float)1/(float)33), 0.0);
	cvConvertScale(dist_temp.R, Dist_ImgStats.Sigma.R, ((float)1/(float)33), 0.0);

	/* Clean up */
	cvReleaseImage(&ref_frame);
	cvReleaseImage(&ref_float);

	cvReleaseImage(&ref_temp.H);
	cvReleaseImage(&ref_temp.V);
	cvReleaseImage(&ref_temp.L);
	cvReleaseImage(&ref_temp.R);

	cvReleaseImage(&dist_frame);
	cvReleaseImage(&dist_float);

	cvReleaseImage(&dist_temp.H);
	cvReleaseImage(&dist_temp.V);
	cvReleaseImage(&dist_temp.L);
	cvReleaseImage(&dist_temp.R);

	return;
}

/*
 * This function calculates the multiplication of the reference and the distorted
 * frames for the calculation of the co-variation between the two
 */
static void calcMul(const IplImage *ref, const IplImage *dist,
					   const vssim_FilStats_t *ref_Mu, const vssim_FilStats_t *dist_Mu,
					   vssim_size_t size, vssim_FilStats_t *temp)
{
	IplImage *ref_sub = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *dist_sub = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *mul = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	IplImage *filH = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *filV = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *filL = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *filR = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	// difference the reference and the Mu for one of the orientations
	cvSub(ref, ref_Mu->H, ref_sub);
	// difference the distorted and the Mu for one of the orientations
	cvSub(dist, dist_Mu->H, dist_sub);
	// multiple the above two differences
	cvMul(ref_sub, dist_sub, mul);
	// apply the corresponding filter
	filterImage(mul, Filters.H, filH);
	// accumulate
	cvAdd(filH, temp->H, temp->H);

	cvSub(ref, ref_Mu->V, ref_sub);
	cvSub(dist, dist_Mu->V, dist_sub);
	cvMul(ref_sub, dist_sub, mul);
	filterImage(mul, Filters.V, filV);
	cvAdd(filH, temp->V, temp->V);

	cvSub(ref, ref_Mu->L, ref_sub);
	cvSub(dist, dist_Mu->L, dist_sub);
	cvMul(ref_sub, dist_sub, mul);
	filterImage(mul, Filters.L, filL);
	cvAdd(filH, temp->L, temp->L);

	cvSub(ref, ref_Mu->R, ref_sub);
	cvSub(dist, dist_Mu->R, dist_sub);
	cvMul(ref_sub, dist_sub, mul);
	filterImage(mul, Filters.R, filR);
	cvAdd(filH, temp->R, temp->R);

	/* Clean up */
	cvReleaseImage(&ref_sub);
	cvReleaseImage(&dist_sub);
	cvReleaseImage(&mul);
	cvReleaseImage(&filH);
	cvReleaseImage(&filV);
	cvReleaseImage(&filL);
	cvReleaseImage(&filR);

	return;
}

/*
 * This function calculates the Co-Variance over the chunk of 33 images
 */
static void calcCovar(FILE *ref, FILE *dist, int cntframe, vssim_size_t size, vssim_ImgComp_t comp)
{
	int i;
	IplImage *temp;

	/* Get 33 frames, 16 previous and 16 future from the file and add them up
	 * separately for Y, U and V */
	IplImage *ref_frame = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_8U, 1);
	IplImage *ref_float =cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	IplImage *dist_frame = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_8U, 1);
	IplImage *dist_float =cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	vssim_FilStats_t stats;
	stats.H = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	stats.V = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	stats.L = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	stats.R = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	long offset;
	switch (comp)
	{
		case Y:
			offset = 0;
			break;
		case U:
			offset = global.size.width * HEIGHT_Y(&global);
			break;
		case V:
			offset = global.size.width *(HEIGHT_Y(&global)+HEIGHT_U(&global));
			break;
	}

	for (i = (cntframe-16); i <= (cntframe+16); i++)
	{
		/* go to the ith frame in the file and read the buffer */
		fseek(ref, ((i*global.size.width*FRAME_HEIGHT(&global))+offset), SEEK_SET);
		fread(ref_frame->imageData, sizeof(char),
				(size.width * size.height),
				ref);
		temp = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
		cvConvert(ref_frame, temp);
		cvResize(temp, ref_float);
		cvReleaseImage(&temp);

		fseek(dist, ((i*global.size.width*FRAME_HEIGHT(&global))+offset), SEEK_SET);
		fread(dist_frame->imageData, sizeof(char),
				(size.width * size.height),
				dist);
		temp = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
		cvConvert(dist_frame, temp);
		cvResize(temp, dist_float);

		/* Subtract H/V/L/R Mu from the images and then Multiply the corresponding
		 * images together to get the variance between the ref and distorted frames */
		calcMul(ref_float, dist_float, &Ref_ImgStats.Mu, &Dist_ImgStats.Mu, size, &stats);

	}

	// div the accumulated value by the num of frames in the chunk, 33 in this case
	cvConvertScale(stats.H, CovarSig.H, ((float)1/(float)33), 0.0);
	cvConvertScale(stats.V, CovarSig.V, ((float)1/(float)33), 0.0);
	cvConvertScale(stats.L, CovarSig.L, ((float)1/(float)33), 0.0);
	cvConvertScale(stats.R, CovarSig.R, ((float)1/(float)33), 0.0);

	/* Clean up */
	cvReleaseImage(&ref_frame);
	cvReleaseImage(&ref_float);

	cvReleaseImage(&dist_frame);
	cvReleaseImage(&dist_float);

	cvReleaseImage(&stats.H);
	cvReleaseImage(&stats.V);
	cvReleaseImage(&stats.L);
	cvReleaseImage(&stats.R);

	return;
}

/*
 * This function calculates the 3D SSIM for the frame in difference orientations
 * using the equation in the paper
 */
static void calc3DSSIM(const IplImage *Ref_Mu, const IplImage *Dist_Mu,
						const IplImage *Ref_Sigma, const IplImage *Dist_Sigma,
						const IplImage *covar, vssim_size_t size, IplImage *ssim3D)
{
	float C1, C2, K1, K2;
	int L = 256; //dynamic range of the image

	K1 = 0.01;
	K2 = 0.03;
	C1 = (K1*L)*(K1*L);
	C2 = (K2*L)*(K2*L);

	// (2*Mu_x*Mu_y + C1)
	IplImage *mul = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *temp1 = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvMul(Ref_Mu, Dist_Mu, mul, 1);
	cvScale(mul, temp1, 2);
	cvAddS(temp1, cvScalarAll(C1), temp1);

	// (2*Covar+C2)
	IplImage *temp2 = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvScale(covar, temp2, 2);
	cvAddS(temp2, cvScalarAll(C2), temp2);

	// (2*Mu_x*Mu_y + C1) * (2*Covar+C2)
	IplImage *temp3 = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvMul(temp1, temp2, temp3, 1);

	IplImage *ref_musq = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	IplImage *dist_musq = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	cvPow(Ref_Mu, ref_musq, 2);
	cvPow(Dist_Mu, dist_musq, 2);

	// (Mu_x^2+Mu_y^2+C1)
	cvAdd(ref_musq, dist_musq, temp1);
	cvAddS(temp1, cvScalarAll(C1), temp1);

	// (sig_x^2+sig_y^2+C1)
	cvAdd(Ref_Sigma, Dist_Sigma, temp2);
	cvAddS(temp2, cvScalarAll(C2), temp2);

	// (Mu_x^2+Mu_y^2+C1)*(sig_x^2+sig_y^2+C1)
	cvMul(temp1, temp2, temp1, 1);

	//3D-SSIM
	cvDiv(temp3, temp1, ssim3D);

	/* Clean up */
	cvReleaseImage(&mul);
	cvReleaseImage(&temp1);
	cvReleaseImage(&temp2);
	cvReleaseImage(&temp3);
	cvReleaseImage(&ref_musq);
	cvReleaseImage(&dist_musq);

	return;
}

/*
 * This function calculates the score for a pixel location in the frame by checking
 * the distance of the motion vector from the two optimal orientations.
 */
static float getScore(vssim_mv_t mv, float h, float v, float l, float r)
{
	float score;
	float avg = 0.0;

	vssim_mv_t plane1, plane2;
	bool pl_hor;
	bool pl_left;
	float d1, d2;	// distances from the 2 planes;

	// if the horizontal comp is greater than vertical, the MV is closer to
	// Horizontal plane, so plane=horizontal
	if (abs(mv.x) > abs(mv.y))
	{
		plane1.x = mv.x;
		plane1.y = 0;
		pl_hor = true;
	}
	// else it is closer to the vertical plane
	else
	{
		plane1.y = mv.y;
		plane1.x = 0;
		pl_hor = false;
	}

	// if the signs on the MV are opposite, then the MV is closer to left plane
	if (((mv.x > 0) && (mv.y < 0)) ||
		((mv.x < 0) && (mv.y > 0)))
	{
		plane2.x = (abs(mv.x) > abs(mv.y)) ? (mv.x) : (-1*mv.y);
		plane2.y = (-1*plane2.x);
		pl_left = true;
	}
	else if (((mv.x > 0) && (mv.y > 0)) ||
			  ((mv.x < 0) && (mv.y < 0)))
	{
		plane2.x = (abs(mv.x) > abs(mv.y)) ? (mv.x) : (mv.y);
		plane2.y = plane2.x;
		pl_left = false;
	}

	// After finding the filters calculate the distance
	// d = sqrt((xi*xi - xj*xj)+(yi*yi - yj*yj))
	d1 = sqrt((mv.x*mv.x - plane1.x*plane1.x) + (mv.y*mv.y - plane1.y*plane1.y));
	d2 = sqrt((mv.x*mv.x - plane2.x*plane2.x) + (mv.y*mv.y - plane2.y*plane2.y));

	// score corresponds to the plane score which the MV is closest to
	if (d1 < d2)
	{
		if (pl_hor == true)
			score = h;
		else
			score = v;
	}
	else if (d1 > d2)
	{
		if (pl_left == true)
			score = l;
		else
			score = r;
	}
	// if the distance of the MV from both the planes is equal then the
	// score = avg of the 2 scores
	else
	{
		if (pl_hor == true)
			avg += h;
		else
			avg += v;

		if (pl_left == true)
			avg += l;
		else
			avg += r;

		score = avg/2.0;
	}

	return score;
}

/*
 * This function calculates the final temporal ssim-map for a frame by collapsing
 * the values from the SSIM-3D along the different orientations using the motion
 * information
 */
static float calcTframe(vssim_mv_t *mv, vssim_size_t size,
							int block_size, int num_blocks)
{
	float Tframe = 1.0;
	int nblock = 0; 	// indicates the block number

	IplImage *MSSIM = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	int step = MSSIM->widthStep;

	for (int h = 0; h < size.height; h++)
	{
		for (int w = 0; w < (size.width*4); w++)
		{
			/* Calculate the MB number corresponding to the current pix, to get
			 * corresponding motion vector */
			nblock = (w/(block_size*4)) + ((h/block_size) * (size.width/(4*block_size)));
			if (nblock > num_blocks)
			{
				break;
			}

			/* Check the direction of the motion vectors */
			// If mv=0, then the score = avg of all the 4 scores
			if ((mv[nblock].x == 0) && (mv[nblock].y == 0))
			{
				MSSIM->imageData[h*step + w] = (SSIM3D.H->imageData[h*step + w] +
											   SSIM3D.V->imageData[h*step + w] +
											   SSIM3D.L->imageData[h*step + w] +
											   SSIM3D.R->imageData[h*step + w]) / 4.0;
			}
			// Check if the mv is horizontal
			else if ((abs(mv[nblock].x) != 0) && (abs(mv[nblock].y) == 0))
			{
				MSSIM->imageData[h*step + w] = SSIM3D.H->imageData[h*step + w];
			}
			// Check if the mv is vertical
			else if ((abs(mv[nblock].x) == 0) && (abs(mv[nblock].y) != 0))
			{
				MSSIM->imageData[h*step + w] = SSIM3D.V->imageData[h*step + w];
			}
			// Check if the mv is along one of the diagonal planes
			else if (abs(mv[nblock].x) == abs(mv[nblock].y))
			{
				// x = -y is the left plane
				if (mv[nblock].x == (-1*mv[nblock].y))
				{
					MSSIM->imageData[h*step + w] = SSIM3D.L->imageData[h*step + w];
				}
				// x = y is the right plane
				else if (mv[nblock].x == mv[nblock].y)
				{
					MSSIM->imageData[h*step + w] = SSIM3D.R->imageData[h*step + w];
				}
			}
			// check for the distances of vectors from the filter planes and
			// select the score that is nearest
			else
			{
				MSSIM->imageData[h*step + w] = (float)getScore(mv[nblock],
														 (float)SSIM3D.H->imageData[h*step + w],
														 (float)SSIM3D.V->imageData[h*step + w],
														 (float)SSIM3D.L->imageData[h*step + w],
														 (float)SSIM3D.R->imageData[h*step + w]);
			}
		}
	}

	// Calculate the avg of the lowest 6% scores
	Tframe = calcPSSIM(MSSIM);

	/* Clean up */
	cvReleaseImage(&MSSIM);

	return Tframe;
}

/*
 * This function returns the final score for the frame.
 * It calls other functions for calculation of Mu, Sigma, Covar and SSIM3D.
 * Then it calls the function for motion vector calculation. It calls another
 * function that uses these motion information for collapsing the different
 * SSIM3D scores into one. Finally calculates the P-SSIM to get the frame score
 */
static float calcMSSIM(FILE *ref, FILE *dist, int cntframe, vssim_ImgComp_t comp, int index)
{
	vssim_size_t size;
	float Tframe;

	size.width = global.size.width;
	switch (comp)
	{
		case Y:
			size.height = HEIGHT_Y(&global);
			break;
		case U:
			size.height = HEIGHT_U(&global);
			break;
		case V:
			size.height = HEIGHT_V(&global);
			break;
	}

	/* Calculate Mu for reference and distorted frames */
	// allocate the required frames
	Ref_ImgStats.Mu.H = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	Ref_ImgStats.Mu.V = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	Ref_ImgStats.Mu.L = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	Ref_ImgStats.Mu.R = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	Dist_ImgStats.Mu.H = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	Dist_ImgStats.Mu.V = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	Dist_ImgStats.Mu.L = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	Dist_ImgStats.Mu.R = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	calcMu(ref, dist, cntframe, size, comp);

	/* Calculate the Standard Deviation for the reference and the distorted frames */
	Ref_ImgStats.Sigma.H = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	Ref_ImgStats.Sigma.V = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	Ref_ImgStats.Sigma.L = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	Ref_ImgStats.Sigma.R = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	Dist_ImgStats.Sigma.H = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	Dist_ImgStats.Sigma.V = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	Dist_ImgStats.Sigma.L = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	Dist_ImgStats.Sigma.R = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	calcSigma(ref, dist, cntframe, size, comp);

	/* Calculate the covariance matrix between the reference and distorted frames */
	CovarSig.H = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	CovarSig.V = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	CovarSig.L = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	CovarSig.R = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	calcCovar(ref, dist, cntframe, size, comp);

	SSIM3D.H = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	SSIM3D.V = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	SSIM3D.L = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);
	SSIM3D.R = cvCreateImage(cvSize(size.width, size.height), IPL_DEPTH_32F, 1);

	/* Calculate the SSIM 3D values in H, V, L and R orientations */
	calc3DSSIM(Ref_ImgStats.Mu.H, Dist_ImgStats.Mu.H,
										  Ref_ImgStats.Sigma.H, Dist_ImgStats.Sigma.H,
										  CovarSig.H, size, SSIM3D.H);

	calc3DSSIM(Ref_ImgStats.Mu.V, Dist_ImgStats.Mu.V,
										  Ref_ImgStats.Sigma.V, Dist_ImgStats.Sigma.V,
										  CovarSig.V, size, SSIM3D.V);

	calc3DSSIM(Ref_ImgStats.Mu.L, Dist_ImgStats.Mu.L,
										  Ref_ImgStats.Sigma.L, Dist_ImgStats.Sigma.L,
										  CovarSig.L, size, SSIM3D.L);

	calc3DSSIM(Ref_ImgStats.Mu.R, Dist_ImgStats.Mu.R,
										  Ref_ImgStats.Sigma.R, Dist_ImgStats.Sigma.R,
										  CovarSig.R, size, SSIM3D.R);

	/* Calculate the Motion vectors for the current image */
	// size of the macro block used
	int block_size = 8;

	// total # of blocks in the image
	int num_blocks = (size.width * size.height)/(block_size * block_size);

	//Stores the direction of MV in X and Y planes.
	vssim_mv_t *mv = (vssim_mv_t *)malloc(num_blocks * sizeof(vssim_mv_t));

	calcMotionVectors(ref, cntframe, size, comp, mv, block_size, num_blocks);

	/* Collapse the 3DSSIM values into one MSSIM score and find out the avg over
	 * the lowest 6% values from the set */
	Tframe = calcTframe(mv, size, block_size, num_blocks);

	/* Clean up */
	free(mv);
	cvReleaseImage(&Ref_ImgStats.Mu.H);
	cvReleaseImage(&Ref_ImgStats.Mu.V);
	cvReleaseImage(&Ref_ImgStats.Mu.L);
	cvReleaseImage(&Ref_ImgStats.Mu.R);

	cvReleaseImage(&Ref_ImgStats.Sigma.H);
	cvReleaseImage(&Ref_ImgStats.Sigma.V);
	cvReleaseImage(&Ref_ImgStats.Sigma.L);
	cvReleaseImage(&Ref_ImgStats.Sigma.R);

	cvReleaseImage(&Dist_ImgStats.Mu.H);
	cvReleaseImage(&Dist_ImgStats.Mu.V);
	cvReleaseImage(&Dist_ImgStats.Mu.L);
	cvReleaseImage(&Dist_ImgStats.Mu.R);

	cvReleaseImage(&Dist_ImgStats.Sigma.H);
	cvReleaseImage(&Dist_ImgStats.Sigma.V);
	cvReleaseImage(&Dist_ImgStats.Sigma.L);
	cvReleaseImage(&Dist_ImgStats.Sigma.R);

	cvReleaseImage(&CovarSig.H);
	cvReleaseImage(&CovarSig.V);
	cvReleaseImage(&CovarSig.L);
	cvReleaseImage(&CovarSig.R);

	cvReleaseImage(&SSIM3D.H);
	cvReleaseImage(&SSIM3D.V);
	cvReleaseImage(&SSIM3D.L);
	cvReleaseImage(&SSIM3D.R);

	return Tframe;
}

/*
 * This function returns the temporal SSIM for the Y, U, V frames from the sequence
 * and returns the final Tframe score to main
 */
vssim_result_t getMSSIM(int cntframe)
{
	int index = (cntframe/16)-1;

	/* Create the 4 Filters to be used */
	createFilters();

	/* Calculate MSSIM for Y */
//	printf("\tT-Score - ");
	global.Tframe[index].score_Y = calcMSSIM(refVideo.fptr, distortedVideo.fptr, cntframe, Y, index);
//	printf("Y: %.3f\t", global.Tframe[index].score_Y);

	/* Calculate MSSIM for U */
	global.Tframe[index].score_U = calcMSSIM(refVideo.fptr, distortedVideo.fptr, cntframe, U, index);
//	printf("U: %.3f\t", global.Tframe[index].score_U);

	/* Calculate MSSIM for V */
	global.Tframe[index].score_V = calcMSSIM(refVideo.fptr, distortedVideo.fptr, cntframe, V, index);
//	printf("V: %.3f\n", global.Tframe[index].score_V);

	/* clean up */
	cvReleaseMat(&Filters.H);
	cvReleaseMat(&Filters.V);
	cvReleaseMat(&Filters.L);
	cvReleaseMat(&Filters.R);

	return SUCCESS;
}
